from elasticsearch import Elasticsearch
import lxconfig
from datetime import datetime

class ElasticsearchStorage:
    def __init__(self):
        es_config = lxconfig.get('elasticsearch')
        self.host = es_config['host']
        self.index = es_config['index']
        self.username = es_config.get('username')
        self.password = es_config.get('password')
        self.client = None

    def connect(self):
        try:
            self.client = Elasticsearch([self.host], http_auth=(self.username, self.password), verify_certs=False)
            print(f'成功连接到Elasticsearch: {self.host}')
        except Exception as e:
            print(f'连接Elasticsearch失败: {str(e)}')

    def store_message(self, message):
        if self.client:
            try:
                day_index = self.index + datetime.now().strftime('%Y%m%d')
                self.client.index(index=day_index, document=message)
            except Exception as e:
                print(f'存储消息到Elasticsearch失败: {str(e)}')
        else:
            print('未建立连接，无法存储消息')

    def bulk_store_messages(self, messages):
        if self.client:
            try:
                day_index = self.index + datetime.now().strftime('%Y%m%d')
                actions = [
                    {
                        '_index': day_index,
                        '_source': message
                    }
                    for message in messages
                ]
                self.client.bulk(operations=actions)
            except Exception as e:
                print(f'批量存储消息到Elasticsearch失败: {str(e)}')
        else:
            print('未建立连接，无法存储消息')

    def close(self):
        if self.client:
            self.client.close()
            print('已关闭Elasticsearch连接')

